import angr
import json
import os
from tqdm import tqdm
from angr.procedures.stubs.ReturnUnconstrained import ReturnUnconstrained

from src.features.angr_adapter import analyses_cfg


class Find_danger_inst:
    def __init__(self,input_dir,danger_list=[]):
        self.input_dir = input_dir
        self.file_list = [file_name for file_name in os.listdir(self.input_dir)]
        danger_juliet = ['wmemset','wcsncpy','memset','fgets','fscanf','memcpy','wcscpy','strcpy','strncpy','strncat',
                         'SNPRINTF','strcat','wcscat','socket','strlen','wcsncat','malloc','memmove']
        self.danger_list = danger_juliet+danger_list

    def analyses_cfg(self, binary_path, min_addr, max_addr):
        project = angr.Project(binary_path, auto_load_libs=False)
        project.hook(0x7001f8, ReturnUnconstrained())
        function_addrs = set()
        unique_addrs = set()
        newly_discovered_addrs = set([project.entry])
        call_depth = 100
        cfg = None
        temp_cfg = None
        while newly_discovered_addrs:
            current_discovered_addrs = set()
            for addr in newly_discovered_addrs:
                if min_addr <= addr <= max_addr and addr not in unique_addrs:
                    temp_cfg = project.analyses.CFGEmulated(
                        context_sensitivity_level=2,
                        keep_state=True,
                        call_depth=call_depth,
                        starts=[addr],
                    )
                    unique_addrs.add(addr)
                    if cfg is None:
                        cfg = temp_cfg
                    else:
                        for func_addr, func in temp_cfg.kb.functions.items():
                            if func_addr not in cfg.kb.functions:
                                cfg.kb.functions[func_addr] = func
                    for func in temp_cfg.kb.functions.values():
                        if func.addr not in function_addrs:
                            current_discovered_addrs.add(func.addr)

            function_addrs.update(newly_discovered_addrs)

            newly_discovered_addrs = current_discovered_addrs - function_addrs

            if not newly_discovered_addrs:
                break
        return cfg

    def find_danger_inst(self):
        all_outputs = []
        for file_name in tqdm(self.file_list, desc="Processing files", unit="file"):
            input_path = os.path.join(self.input_dir, file_name)
            proj = angr.Project(input_path, auto_load_libs=False)
            main_obj = proj.loader.main_object
            min_addr = main_obj.min_addr
            max_addr = main_obj.max_addr
            cfg = self.analyses_cfg(input_path, min_addr, max_addr)
            plt_table = proj.loader.main_object.plt
            dangerous_functions = [entry for entry in plt_table if entry in self.danger_list]
            all_ins = []
            for idx,item in enumerate(dangerous_functions):
                danger_addr = plt_table[item]
                    # print(f"memcpy found at PLT address: {hex(memcpy_addr)}")
                func_node = cfg.get_any_node(danger_addr)
                successors_1 = cfg.get_successors(func_node)
                if successors_1:
                    successors_2 = cfg.get_successors(successors_1[0])
                    if successors_2:
                        successors_2=successors_2[0]
                        # print(f"Successor address: {hex(successors_2.addr)}")
                        block = proj.factory.block(successors_2.addr)
                        # print(f"Instructions in block at {hex(block.addr)}:")
                        all_ins.append({"idx":idx})
                        for insn in block.capstone.insns:
                            all_ins.append({
                                    'addr': insn.address,
                                    'nemonic': insn.mnemonic,
                                    'op_str': insn.op_str
                                    })
                        if all_ins:
                                all_ins = all_ins[:-1]
                        print(all_ins)
                else:
                    all_ins.append({"idx":idx})
            all_outputs.append({
                'file': file_name,
                'instructions': all_ins,
                'dangerous_functions':dangerous_functions
            })
        output_dir = os.path.basename(self.input_dir)
        output_file_name = os.path.join(output_dir + '.json')
        with open(output_file_name, 'a') as json_file:
            json.dump(all_outputs, json_file)

if __name__ == '__main__':
    a=Find_danger_inst('/home/wang1/zz_static_project/s01/s02')
    a.find_danger_inst()
